from torch.utils.data import Dataset
import torch
import numpy as np

'''
    注意：这是数据集的类，输入的result是readData   
'''
class VoxedgeDataset(Dataset):

    '''
        说明：voxdata：仅仅是继承了Dataset类的一个具体的对象，
                      下述方法中调用的使用get__items__(),
                      获取对应idx对应的三维坐标，和当前点是否是边界点判定条件
             edges
        参数：输入的数据是修改过之后，是包含所有的采样点，但是仅仅只有表面点被标注为1
             result形如[[x,y,z,1](该点是表面点),[x,y,z,0]（该点不是表面点）]
             遍历了体素空间中的所有的点，128*128*128个点
    '''
    def __init__(self, result):
        self.data = []
        # 统计体素表面对应的点
        self.datalen = 0
        # 添加的数据是张良对应的三维坐标tensor类型的三维坐标tensor[x,y,z]
        # 和该点是否存在对应点的tensor的long型的数据
        for i in result:
            a = np.array(i).astype(np.int)
            self.data.append((torch.Tensor([a[0]/128,a[1]/128,a[2]/128,a[3]])))
            if(i[3] == 1):
                self.datalen = self.datalen + 1

    # 获取一维坐标idx对应的三维体素坐标点以及对应点是否有体素的情况
    def __getitem__(self, idx):
        return self.data[idx]

    # 统计模型表面点的个数
    def __len__(self):
        return self.datalen

    # 删除表面的点
    def remove(self,a,b,c,d):
        for i in self.data:
            if(a*i[0]+b*i[1]+c*i[2]+d*i[3] > -1 and a*i[0]+b*i[1]+c*i[2]+d*i[3] < 1):
                self.remove(a)
